from typing import List, Dict
import json
import random
from http import HTTPStatus
from io import BytesIO
from urllib.request import urlopen
import os

from PIL import Image
from dashscope import ImageSynthesis

from mm_story_agent.prompts_en import role_extract_system, role_review_system, \
    story_to_image_reviser_system, story_to_image_review_system
from mm_story_agent.base import register_tool, init_tool_instance
from mm_story_agent.utils.long_text import LongTextProcessor


class ImageAPISynthesizer:

    def __init__(self,
                 num_pages: int,
                 height: int,
                 width: int,
                 model_name: str = "wanx-v1",
                 num_steps: int = 30):
        self.total_length = num_pages
        self.height = height
        self.width = width
        self.model_name = model_name
        self.num_steps = num_steps

        # Prompt styles retained to keep backward compatibility with previous UX
        self.styles = {
            '(No style)': (
                '{prompt}',
                ''),
            'Japanese Anime': (
                'anime artwork illustrating {prompt}. created by japanese anime studio. highly emotional. best quality, high resolution, (Anime Style, Manga Style:1.3), Low detail, sketch, concept art, line art, webtoon, manhua, hand drawn, defined lines, simple shades, minimalistic, High contrast, Linear compositions, Scalable artwork, Digital art, High Contrast Shadows',
                'lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'),
            'Digital/Oil Painting': (
                '{prompt} . (Extremely Detailed Oil Painting:1.2), glow effects, godrays, hand drawn, render, ultra detailed, cinematic, depth of field',
                'anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'),
            'Pixar/Disney Character': (
                'Create a Disney Pixar 3D style illustration on {prompt}. Vibrant, motivational, vivid colors, a sense of wonder.',
                'lowres, bad anatomy, bad hands, text, bad eyes, bad arms, bad legs, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, blurry, grayscale, noisy, sloppy, messy, grainy, ultra textured, photo'),
            'Photographic': (
                'cinematic photo of {prompt}. Hyperrealistic, detailed skin, soft lighting, realistic, best quality, ultra realistic, film photography, soft focus',
                'drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'),
            'Comic book': (
                'comic illustration of {prompt}. graphic novel art, vibrant, highly detailed',
                'photograph, deformed, glitch, noisy, realistic, stock photo, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry'),
            'Line art': (
                'line art drawing of {prompt}. professional, sleek, modern, minimalist, vector graphics',
                'photorealistic, 35mm film, deformed, glitch, blurry, noisy, off-center, cross-eyed, bad anatomy, ugly, disfigured, mutated, realism, oil, acrylic, lowres, bad hands, text, error, missing fingers, extra digit, cropped, jpeg artifacts, watermark'),
            'Black and White Film Noir': (
                '{prompt}. Monochromatic film photography, film noir, analog style, soft lighting, heavy shadow, masterpiece',
                'anime, photorealistic color, deformed, glitch, blurry, noisy, off-center, bad anatomy, mutated, lowres, bad hands, text, error, jpeg artifacts, watermark'),
            'Isometric Rooms': (
                'Tiny cute isometric {prompt} in a cutaway box, soft smooth lighting, soft colors, 100mm lens, 3d blender render',
                'photorealistic, deformed, glitch, blurry, noisy, bad anatomy, text, error, cropped, jpeg artifacts, watermark'),
            'Storybook': (
                'Cartoon style, cute illustration of {prompt}.',
                'realism, photo, realistic, lowres, bad hands, bad eyes, bad arms, bad legs, error, missing fingers, cropped, worst quality, jpeg artifacts, signature, grayscale, noisy, messy, grainy, ultra textured'
            )
        }

        self.default_negative = (
            'naked, deformed, bad anatomy, disfigured, poorly drawn face, mutation, '
            'extra limb, ugly, disgusting, poorly drawn hands, missing limb, floating limbs, '
            'disconnected limbs, blurry, watermarks, oversaturated, distorted hands, amputation'
        )

    def apply_style(self, style_name: str, positives: List[str], negative: str = ""):
        p, n = self.styles.get(style_name, self.styles['(No style)'])
        pos = [p.replace('{prompt}', s) for s in positives]
        neg = (n + ' ' + negative).strip()
        return pos, neg

    def apply_style_positive(self, style_name: str, positive: str):
        p, _ = self.styles.get(style_name, self.styles['(No style)'])
        return p.replace('{prompt}', positive)

    def _synthesize_one(self, prompt: str, negative_prompt: str, seed: int, max_try: int = 3):
        try_id = 0
        last_err = None
        while try_id < max_try:
            try:
                result = ImageSynthesis.call(
                    model=self.model_name,
                    prompt=prompt,
                    negative_prompt=negative_prompt or None,
                    size=f"{self.width}*{self.height}",
                    n=1,
                    seed=seed,
                )
                if result.status_code == HTTPStatus.OK and result.output and result.output.get('results'):
                    url = result.output['results'][0].get('url')
                    if not url:
                        raise RuntimeError('Empty image url in result.')
                    with urlopen(url, timeout=120) as resp:
                        data = resp.read()
                    img = Image.open(BytesIO(data)).convert('RGB')
                    return img
                else:
                    last_err = RuntimeError(f"Image API failed: status={result.status_code}, code={getattr(result, 'code', None)}, message={getattr(result, 'message', None)}")
            except Exception as e:
                last_err = e
            try_id += 1
            seed = (seed + 1) % 1000000007  # minor jitter for retry
        raise last_err if last_err else RuntimeError('Unknown image synthesis error.')

    def call(self,
             prompts: List[str],
             style_name: str = 'Storybook',
             guidance_scale: float = 5.0,  # kept for interface compatibility, not used by API
             seed: int = 2047):
        assert len(prompts) == self.total_length, "The number of prompts should be equal to the number of pages."

        # Apply style once for the first id images and then per image for the rest (kept behavior shape)
        id_length = 0  # API does not support ID-consistency; generate per prompt directly
        images: List[Image.Image] = []

        # Build negative once
        _, negative_prompt = self.apply_style(style_name, [], self.default_negative)

        for idx, p in enumerate(prompts):
            styled_prompt = self.apply_style_positive(style_name, p)
            img = self._synthesize_one(styled_prompt, negative_prompt, seed=seed + idx)
            images.append(img)

        return images


@register_tool("story_diffusion_t2i")
class StoryDiffusionAgent:

    def __init__(self, cfg) -> None:
        self.cfg = cfg
        
    def call(self, params: Dict):
        pages: List = params["pages"]
        save_path: str = params["save_path"]
        # 长文本压缩（若有必要）
        pages = LongTextProcessor().compact_pages_if_needed(pages)
        role_dict = self.extract_role_from_story(pages)
        image_prompts = self.generate_image_prompt_from_story(pages)
        image_prompts_with_role_desc = []
        for image_prompt in image_prompts:
            for role, role_desc in role_dict.items():
                if role in image_prompt:
                    image_prompt = image_prompt.replace(role, role_desc)
            image_prompts_with_role_desc.append(image_prompt)
        # 模型选择优先级：params.model_name > cfg.model_name > 环境变量 DASHSCOPE_IMAGE_MODEL > 默认 wanx-v1
        chosen_model = (
            params.get("model_name")
            or self.cfg.get("model_name")
            or os.environ.get("DASHSCOPE_IMAGE_MODEL")
            or "wanx-v1"
        )
        generation_agent = ImageAPISynthesizer(
            num_pages=len(pages),
            height=self.cfg.get("height", 512),
            width=self.cfg.get("width", 512),
            model_name=chosen_model,
            num_steps=self.cfg.get("num_steps", 30)
        )
        images = generation_agent.call(
            image_prompts_with_role_desc,
            style_name=params.get("style_name", "Storybook"),
            guidance_scale=params.get("guidance_scale", 5.0),
            seed=params.get("seed", 2047)
        )
        for idx, image in enumerate(images):
            image.save(save_path / f"p{idx + 1}.png")
        return {
            "prompts": image_prompts_with_role_desc,
            "generation_results": images,
        }
        
    def extract_role_from_story(
            self,
            pages: List,
        ):
        num_turns = self.cfg.get("num_turns", 3)
        role_extractor = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": {
                "system_prompt": role_extract_system,
                "track_history": False
            }
        })
        role_reviewer = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": {
                "system_prompt": role_review_system,
                "track_history": False
            }
        })
        roles = {}
        review = ""
        for turn in range(num_turns):
            roles, success = role_extractor.call(json.dumps({
                    "story_content": pages,
                    "previous_result": roles,
                    "improvement_suggestions": review,
                }, ensure_ascii=False
            ))
            roles = json.loads(roles.strip("```json").strip("```"))
            review, success = role_reviewer.call(json.dumps({
                "story_content": pages,
                "role_descriptions": roles
            }, ensure_ascii=False))
            if review == "Check passed.":
                break
        return roles

    def generate_image_prompt_from_story(
            self,
            pages: List,
            num_turns: int = 3
        ):
        image_prompt_reviewer = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": {
                "system_prompt": story_to_image_review_system,
                "track_history": False
            }
        })
        image_prompt_reviser = init_tool_instance({
            "tool": self.cfg.get("llm", "qwen"),
            "cfg": {
                "system_prompt": story_to_image_reviser_system,
                "track_history": False
            }
        })
        image_prompts = []

        for page in pages:
            review = ""
            image_prompt = ""
            for turn in range(num_turns):
                image_prompt, success = image_prompt_reviser.call(json.dumps({
                    "all_pages": pages,
                    "current_page": page,
                    "previous_result": image_prompt,
                    "improvement_suggestions": review,
                }, ensure_ascii=False))
                if image_prompt.startswith("Image description:"):
                    image_prompt = image_prompt[len("Image description:"):]
                review, success = image_prompt_reviewer.call(json.dumps({
                    "all_pages": pages,
                    "current_page": page,
                    "image_description": image_prompt
                }, ensure_ascii=False))
                if review == "Check passed.":
                    break
            image_prompts.append(image_prompt)
        return image_prompts

